local component = require("component")
local computer = require("computer")
local event = require("event")
local thread = require("thread")
local serialization = require("serialization")
local filesystem = require("filesystem")

-- ========== 配置区域 ==========
local DHCP_SERVER_PORT = 67
local DHCP_CLIENT_PORT = 68
local ARP_PORT = 69
local ICMP_PORT = 70
local CONFIG_DIR = "/etc/dhcp/"
local DHCP_CHECK_INTERVAL = 300
local ARP_CACHE_TIME = 300
local PING_TIMEOUT = 3
local ARP_TIMEOUT = 2
-- =============================

-- 全局变量
local modem = nil
local mode = nil -- "server" or "client"
local running = true
local threads = {}

-- 网络状态
local networkConfig = {
    ip = nil,
    netmask = "255.255.255.0",
    gateway = nil,
    dns = nil,
    lease_time = 0,
    dhcp_enabled = true
}

-- 服务器配置
local serverConfig = {
    server_ip = "192.168.1.1",
    ip_pool = {
        "192.168.1.100", "192.168.1.101", "192.168.1.102",
        "192.168.1.103", "192.168.1.104", "192.168.1.105"
    },
    lease_time = 3600
}

-- 共享数据结构
local arpTable = {}    -- {[IP] = {mac, timestamp}}
local dhcpLeases = {}  -- {[MAC] = {ip, expire_time, sender}} (仅服务器模式)
local statistics = {
    dhcp_requests = 0,
    dhcp_offers = 0,
    dhcp_acks = 0,
    arp_requests = 0,
    arp_replies = 0,
    icmp_requests = 0,
    icmp_replies = 0,
    start_time = 0
}

-- ===============================
-- 工具函数
-- ===============================

local function selectNetworkCard()
    local modems = {}
    for address in component.list("modem") do
        table.insert(modems, {address = address, proxy = component.proxy(address)})
    end

    if #modems == 0 then
        print("Error: No network card found!")
        return nil
    elseif #modems == 1 then
        print("Using network card: " .. modems[1].address:sub(1,8) .. "...")
        return modems[1].proxy
    else
        print("Multiple network cards detected:")
        for i, modem in ipairs(modems) do
            local isWireless = modem.isWireless and modem.isWireless() or false
            local cardType = isWireless and "(Wireless)" or "(Wired)"
            print(string.format("%d. %s %s", i, modem.address, cardType))
        end

        io.write("Enter card number (1-" .. #modems .. "): ")
        local choice = tonumber(io.read())

        if choice and choice >= 1 and choice <= #modems then
            print("Selected: " .. modems[choice].address:sub(1,8) .. "...")
            return modems[choice].proxy
        else
            print("Invalid selection, using first card")
            return modems[1].proxy
        end
    end
end

local function loadConfig()
    if not filesystem.exists(CONFIG_DIR) then
        filesystem.makeDirectory(CONFIG_DIR)
    end

    local configFile = CONFIG_DIR .. mode .. ".cfg"
    if filesystem.exists(configFile) then
        local file = io.open(configFile, "r")
        if file then
            local content = file:read("*a")
            file:close()
            local success, config = pcall(serialization.unserialize, content)
            if success and config then
                if mode == "client" then
                    for k, v in pairs(config.network or {}) do
                        networkConfig[k] = v
                    end
                elseif mode == "server" then
                    for k, v in pairs(config.server or {}) do
                        serverConfig[k] = v
                    end
                    dhcpLeases = config.leases or {}
                end
                arpTable = config.arp or {}
                statistics = config.stats or statistics
                return true
            end
        end
    end
    return false
end

local function saveConfig()
    local configFile = CONFIG_DIR .. mode .. ".cfg"
    local config = {
        arp = arpTable,
        stats = statistics
    }

    if mode == "client" then
        config.network = networkConfig
    elseif mode == "server" then
        config.server = serverConfig
        config.leases = dhcpLeases
    end

    local file = io.open(configFile, "w")
    if file then
        file:write(serialization.serialize(config))
        file:close()
        return true
    end
    return false
end

local function cleanupArpTable()
    local now = computer.uptime()
    for ip, entry in pairs(arpTable) do
        if now - entry.timestamp > ARP_CACHE_TIME then
            arpTable[ip] = nil
        end
    end
end

local function isInSameSubnet(ip1, ip2, netmask)
    local subnet1 = ip1:match("(%d+%.%d+%.%d+%.)%d+")
    local subnet2 = ip2:match("(%d+%.%d+%.%d+%.)%d+")
    return subnet1 == subnet2
end

local function getMyIP()
    return mode == "server" and serverConfig.server_ip or networkConfig.ip
end

local function hasValidIP()
    if mode == "server" then
        return serverConfig.server_ip ~= nil
    else
        return networkConfig.ip and networkConfig.lease_time > computer.uptime()
    end
end

-- ===============================
-- ARP 协议实现 (共享)
-- ===============================

local function arpRequest(targetIP)
    local myIP = getMyIP()
    if not myIP then return nil end

    cleanupArpTable()

    -- 检查ARP缓存
    if arpTable[targetIP] and (computer.uptime() - arpTable[targetIP].timestamp < ARP_CACHE_TIME) then
        return arpTable[targetIP].mac
    end

    print("[ARP] Resolving " .. targetIP .. "...")

    modem.broadcast(ARP_PORT, serialization.serialize({
        type = "arp_request",
        sender_ip = myIP,
        sender_mac = modem.address,
        target_ip = targetIP
    }))

    local deadline = computer.uptime() + ARP_TIMEOUT
    while computer.uptime() < deadline do
        local eventData = {event.pull(0.1, "modem_message")}
        if eventData[1] == "modem_message" then
            local _, _, from, port, _, payload = table.unpack(eventData)

            if port == ARP_PORT and payload then
                local success, packet = pcall(serialization.unserialize, payload)
                if success and packet and packet.type == "arp_reply" and
                        packet.target_ip == myIP and packet.sender_ip == targetIP then

                    arpTable[targetIP] = {
                        mac = packet.sender_mac,
                        timestamp = computer.uptime()
                    }

                    print("[ARP] " .. targetIP .. " is at " .. packet.sender_mac:sub(1,8) .. "...")
                    return packet.sender_mac
                end
            end
        end
    end

    print("[ARP] No response from " .. targetIP)
    return nil
end

local function handleArpPacket(sender, packet)
    local myIP = getMyIP()
    if not myIP then return end

    if packet.type == "arp_request" then
        statistics.arp_requests = statistics.arp_requests + 1

        -- 更新发送者信息到ARP表
        arpTable[packet.sender_ip] = {
            mac = packet.sender_mac,
            timestamp = computer.uptime()
        }

        print("[ARP] ARP request from " .. packet.sender_ip ..
                " (" .. packet.sender_mac:sub(1,8) .. "...) looking for " .. packet.target_ip)

        -- 如果目标IP是自己
        if packet.target_ip == myIP then
            modem.send(packet.sender_mac, ARP_PORT, serialization.serialize({
                type = "arp_reply",
                sender_ip = myIP,
                sender_mac = modem.address,
                target_ip = packet.sender_ip,
                target_mac = packet.sender_mac
            }))
            statistics.arp_replies = statistics.arp_replies + 1
            print("[ARP] Replied: " .. myIP .. " is at " .. modem.address:sub(1,8) .. "...")

            -- 如果是服务器模式，可以转发ARP请求
        elseif mode == "server" then
            for mac, lease in pairs(dhcpLeases) do
                if mac ~= "server" and lease.ip == packet.target_ip and lease.expire > computer.uptime() then
                    print("[ARP] Forwarding ARP request to " .. packet.target_ip)
                    modem.send(mac, ARP_PORT, serialization.serialize(packet))
                    break
                end
            end
        end

    elseif packet.type == "arp_reply" then
        -- 更新ARP表
        arpTable[packet.sender_ip] = {
            mac = packet.sender_mac,
            timestamp = computer.uptime()
        }

        -- 如果回复的目标不是自己且是服务器模式，转发回复
        if packet.target_ip ~= myIP and mode == "server" then
            for mac, lease in pairs(dhcpLeases) do
                if lease.ip == packet.target_ip then
                    print("[ARP] Forwarding ARP reply from " .. packet.sender_ip .. " to " .. packet.target_ip)
                    modem.send(mac, ARP_PORT, serialization.serialize(packet))
                    break
                end
            end
        end
    end
end

-- ===============================
-- ICMP 协议实现 (共享)
-- ===============================

local function ping(targetIP, count)
    local myIP = getMyIP()
    if not myIP then
        print("Error: No IP address configured")
        return
    end

    count = count or 4
    local sent = 0
    local received = 0
    local responses = {}

    print("PING " .. targetIP .. " from " .. myIP)

    -- 进行ARP解析
    local targetMAC = nil

    if isInSameSubnet(myIP, targetIP, networkConfig.netmask) then
        targetMAC = arpRequest(targetIP)
    elseif networkConfig.gateway then
        print("Target is remote, using gateway " .. networkConfig.gateway)
        targetMAC = arpRequest(networkConfig.gateway)
    else
        print("No gateway configured for remote target")
        return
    end

    if not targetMAC then
        print("ARP failed - cannot reach " .. targetIP)
        return
    end

    for i = 1, count do
        local startTime = computer.uptime()
        local pingId = math.random(1, 65535)

        modem.send(targetMAC, ICMP_PORT, serialization.serialize({
            type = "icmp_echo_request",
            target_ip = targetIP,
            source_ip = myIP,
            source_mac = modem.address,
            id = pingId,
            seq = i,
            timestamp = startTime
        }))

        sent = sent + 1
        local replied = false

        local deadline = startTime + PING_TIMEOUT
        while computer.uptime() < deadline and not replied do
            local eventData = {event.pull(0.1, "modem_message")}
            if eventData[1] == "modem_message" then
                local _, _, from, port, _, payload = table.unpack(eventData)
                if port == ICMP_PORT and payload then
                    local success, packet = pcall(serialization.unserialize, payload)
                    if success and packet and packet.id == pingId and
                            packet.source_ip == targetIP and packet.target_ip == myIP then

                        if packet.type == "icmp_echo_reply" then
                            local responseTime = (computer.uptime() - startTime) * 1000
                            table.insert(responses, responseTime)
                            print(string.format("Reply from %s: seq=%d time=%.1fms",
                                    targetIP, packet.seq, responseTime))
                            received = received + 1
                            replied = true
                        end
                    end
                end
            end
        end

        if not replied then
            print("Request timeout for seq=" .. i)
        end

        if i < count then os.sleep(1) end
    end

    -- 统计信息
    print("\n--- " .. targetIP .. " ping statistics ---")
    print(string.format("%d packets transmitted, %d received, %.1f%% packet loss",
            sent, received, (sent - received) / sent * 100))

    if #responses > 0 then
        local total = 0
        local min_time, max_time = responses[1], responses[1]

        for _, time in ipairs(responses) do
            total = total + time
            if time < min_time then min_time = time end
            if time > max_time then max_time = time end
        end

        local avg_time = total / #responses
        print(string.format("round-trip min/avg/max = %.1f/%.1f/%.1f ms",
                min_time, avg_time, max_time))
    end
end

local function handleIcmpPacket(sender, packet)
    local myIP = getMyIP()
    if not myIP then return end

    if packet.type == "icmp_echo_request" then
        statistics.icmp_requests = statistics.icmp_requests + 1

        print("[ICMP] Ping from " .. packet.source_ip .. " to " .. packet.target_ip ..
                " (seq=" .. (packet.seq or "?") .. ")")

        if packet.target_ip == myIP then
            modem.send(packet.source_mac, ICMP_PORT, serialization.serialize({
                type = "icmp_echo_reply",
                target_ip = packet.source_ip,
                source_ip = myIP,
                source_mac = modem.address,
                target_mac = packet.source_mac,
                id = packet.id,
                seq = packet.seq,
                timestamp = packet.timestamp
            }))
            statistics.icmp_replies = statistics.icmp_replies + 1
            print("[ICMP] Replied to ping from " .. packet.source_ip)

        elseif mode == "server" then
            -- 服务器转发ping包
            for mac, lease in pairs(dhcpLeases) do
                if mac ~= "server" and lease.ip == packet.target_ip and lease.expire > computer.uptime() then
                    print("[ICMP] Forwarding ping from " .. packet.source_ip .. " to " .. packet.target_ip)
                    modem.send(mac, ICMP_PORT, serialization.serialize(packet))
                    return
                end
            end

            print("[ICMP] No route to " .. packet.target_ip)
        end

    elseif packet.type == "icmp_echo_reply" then
        if packet.target_ip ~= myIP and mode == "server" then
            for mac, lease in pairs(dhcpLeases) do
                if lease.ip == packet.target_ip then
                    print("[ICMP] Forwarding ping reply from " .. packet.source_ip .. " to " .. packet.target_ip)
                    modem.send(mac, ICMP_PORT, serialization.serialize(packet))
                    break
                end
            end
        end
    end
end

-- ===============================
-- DHCP 客户端功能
-- ===============================

local function performDHCP()
    if not networkConfig.dhcp_enabled then return false end

    print("[DHCP] Requesting IP address...")

    for attempt = 1, 3 do
        print("[DHCP] Attempt " .. attempt .. "/3")

        modem.broadcast(DHCP_SERVER_PORT, serialization.serialize({
            type = "DISCOVER",
            mac = modem.address,
            client_port = DHCP_CLIENT_PORT
        }))

        local deadline = computer.uptime() + 5
        while computer.uptime() < deadline do
            local eventData = {event.pull(0.1, "modem_message")}
            if eventData[1] == "modem_message" then
                local _, _, from, port, _, payload = table.unpack(eventData)

                if port == DHCP_CLIENT_PORT and payload then
                    local success, packet = pcall(serialization.unserialize, payload)
                    if success and packet and packet.type == "OFFER" and packet.mac == modem.address then
                        print("[DHCP] Received offer: " .. packet.ip)

                        modem.send(from, DHCP_SERVER_PORT, serialization.serialize({
                            type = "REQUEST",
                            mac = modem.address,
                            requested_ip = packet.ip,
                            client_port = DHCP_CLIENT_PORT
                        }))

                        local ack_deadline = computer.uptime() + 3
                        while computer.uptime() < ack_deadline do
                            local ackData = {event.pull(0.1, "modem_message")}
                            if ackData[1] == "modem_message" then
                                local _, _, ack_from, ack_port, _, ack_payload = table.unpack(ackData)
                                if ack_port == DHCP_CLIENT_PORT and ack_payload then
                                    local ack_success, ack_packet = pcall(serialization.unserialize, ack_payload)
                                    if ack_success and ack_packet and ack_packet.type == "ACK" and ack_packet.mac == modem.address then
                                        networkConfig.ip = ack_packet.ip
                                        networkConfig.netmask = packet.netmask or "255.255.255.0"
                                        networkConfig.gateway = packet.gateway
                                        networkConfig.lease_time = computer.uptime() + (packet.lease_time or 3600)

                                        -- 添加网关和服务器到ARP表
                                        if networkConfig.gateway then
                                            arpTable[networkConfig.gateway] = {
                                                mac = from,
                                                timestamp = computer.uptime()
                                            }
                                        end

                                        -- 添加自己的IP到ARP表
                                        arpTable[networkConfig.ip] = {
                                            mac = modem.address,
                                            timestamp = computer.uptime()
                                        }

                                        saveConfig()
                                        print("[DHCP] Success! IP: " .. ack_packet.ip)
                                        return true
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    print("[DHCP] Failed to obtain IP address")
    return false
end

local function dhcpChecker()
    while running and mode == "client" do
        if networkConfig.dhcp_enabled and not hasValidIP() then
            print("[DHCP] IP lease expired or missing, requesting new lease...")
            performDHCP()
        end

        local endTime = computer.uptime() + DHCP_CHECK_INTERVAL
        while computer.uptime() < endTime and running do
            os.sleep(1)
        end
    end
end

-- ===============================
-- DHCP 服务器功能
-- ===============================

local function cleanupExpiredLeases()
    local now = computer.uptime()
    local cleaned = 0
    for mac, lease in pairs(dhcpLeases) do
        if mac ~= "server" and lease.expire < now then
            print("[DHCP] Lease expired for " .. mac:sub(1,8) .. "... (IP: " .. lease.ip .. ")")
            arpTable[lease.ip] = nil
            dhcpLeases[mac] = nil
            cleaned = cleaned + 1
        end
    end
    return cleaned
end

local function handleDhcpPacket(sender, packet)
    if mode ~= "server" then return end

    local mac = packet.mac
    local client_port = packet.client_port or DHCP_CLIENT_PORT

    statistics.dhcp_requests = statistics.dhcp_requests + 1

    if packet.type == "DISCOVER" then
        cleanupExpiredLeases()

        local assignedIp = nil

        if dhcpLeases[mac] and dhcpLeases[mac].expire > computer.uptime() then
            assignedIp = dhcpLeases[mac].ip
            print("[DHCP] Renewing existing lease: " .. assignedIp)
        else
            for _, ip in ipairs(serverConfig.ip_pool) do
                local alreadyAssigned = false
                for leaseMac, lease in pairs(dhcpLeases) do
                    if leaseMac ~= "server" and lease.ip == ip and lease.expire > computer.uptime() then
                        alreadyAssigned = true
                        break
                    end
                end
                if not alreadyAssigned then
                    assignedIp = ip
                    break
                end
            end
        end

        if assignedIp then
            dhcpLeases[mac] = {
                ip = assignedIp,
                expire = computer.uptime() + serverConfig.lease_time,
                sender = sender
            }

            arpTable[assignedIp] = {
                mac = mac,
                timestamp = computer.uptime()
            }

            modem.send(sender, client_port, serialization.serialize({
                type = "OFFER",
                mac = mac,
                ip = assignedIp,
                netmask = "255.255.255.0",
                gateway = serverConfig.server_ip,
                dns = serverConfig.server_ip,
                lease_time = serverConfig.lease_time
            }))
            statistics.dhcp_offers = statistics.dhcp_offers + 1
            print("[DHCP] Offered IP " .. assignedIp .. " to " .. mac:sub(1,8) .. "...")
        end

    elseif packet.type == "REQUEST" then
        if dhcpLeases[mac] and dhcpLeases[mac].ip == packet.requested_ip then
            dhcpLeases[mac].expire = computer.uptime() + serverConfig.lease_time
            dhcpLeases[mac].sender = sender

            arpTable[packet.requested_ip] = {
                mac = mac,
                timestamp = computer.uptime()
            }

            modem.send(sender, client_port, serialization.serialize({
                type = "ACK",
                mac = mac,
                ip = dhcpLeases[mac].ip
            }))
            statistics.dhcp_acks = statistics.dhcp_acks + 1
            print("[DHCP] Confirmed IP " .. dhcpLeases[mac].ip .. " for " .. mac:sub(1,8) .. "...")
        end
    end
end

-- ===============================
-- 网络协议栈监听器
-- ===============================

local function networkListener()
    while running do
        local eventData = {event.pull(0.5, "modem_message")}

        if eventData[1] == "modem_message" then
            local _, _, sender, port, _, payload = table.unpack(eventData)

            if payload then
                local success, packet = pcall(serialization.unserialize, payload)

                if success and packet and packet.type then
                    if port == ARP_PORT then
                        handleArpPacket(sender, packet)
                    elseif port == ICMP_PORT then
                        handleIcmpPacket(sender, packet)
                    elseif port == DHCP_SERVER_PORT and mode == "server" then
                        handleDhcpPacket(sender, packet)
                    end
                end
            end
        end

        -- 定期清理和保存
        if math.floor(computer.uptime()) % 300 == 0 then
            cleanupArpTable()
            if mode == "server" then
                cleanupExpiredLeases()
            end
            saveConfig()
        end
    end
end

-- ===============================
-- CLI 命令处理
-- ===============================

local function showHelp()
    print("Available commands:")
    if mode == "client" then
        print("  ipconfig                  - Show network configuration")
        print("  ipconfig /release         - Release DHCP lease")
        print("  ipconfig /renew           - Renew DHCP lease")
        print("  ipconfig /set <ip>        - Set static IP")
        print("  dhcp <enable|disable>     - Enable/disable DHCP")
    elseif mode == "server" then
        print("  list                      - Show active DHCP leases")
        print("  kick <ip>                 - Kick client (revoke lease)")
    end
    print("  ping <target> [count]     - Ping target address")
    print("  arp                       - Show ARP table")
    print("  arp <ip>                  - Resolve specific IP")
    print("  stats                     - Show statistics")
    print("  save                      - Save configuration")
    print("  reload                    - Reload configuration")
    print("  exit                      - Exit program")
    print("  help                      - Show this help")
end

local function showIPConfig()
    print("\nNetwork Configuration (" .. mode .. " mode):")
    print("======================")
    print("Network Card: " .. modem.address)

    if mode == "client" then
        print("IP Address  : " .. (networkConfig.ip or "Not assigned"))
        print("Subnet Mask : " .. networkConfig.netmask)
        print("Gateway     : " .. (networkConfig.gateway or "Not set"))
        print("DNS Server  : " .. (networkConfig.dns or "Not set"))
        print("DHCP        : " .. (networkConfig.dhcp_enabled and "Enabled" or "Disabled"))

        if networkConfig.dhcp_enabled and networkConfig.lease_time > 0 then
            local remaining = networkConfig.lease_time - computer.uptime()
            if remaining > 0 then
                print("Lease Time  : " .. math.floor(remaining / 60) .. " minutes remaining")
            else
                print("Lease Time  : Expired")
            end
        end
    elseif mode == "server" then
        print("Server IP   : " .. serverConfig.server_ip)
        print("IP Pool     : " .. table.concat(serverConfig.ip_pool, ", "))
        print("Lease Time  : " .. serverConfig.lease_time .. " seconds")
    end
    print("======================\n")
end

local function showArpTable()
    cleanupArpTable()
    print("\nARP Table:")
    print("=========================")
    print("IP Address      MAC Address       Age")
    print("--------------------------------------")

    local now = computer.uptime()
    local count = 0
    for ip, entry in pairs(arpTable) do
        local age = math.floor(now - entry.timestamp)
        print(string.format("%-15s %-17s %d sec", ip, entry.mac:sub(1,17) .. "...", age))
        count = count + 1
    end

    if count == 0 then
        print("No entries")
    end
    print("==========================\n")
end

local function showLeases()
    if mode ~= "server" then
        print("DHCP leases only available in server mode")
        return
    end

    cleanupExpiredLeases()
    print("\n=== Active DHCP Leases ===")
    print("MAC Address     \tIP Address    \tExpires In")
    print("--------------------------------------------")

    local now = computer.uptime()
    local count = 0
    for mac, lease in pairs(dhcpLeases) do
        if mac ~= "server" then
            local remaining = math.max(0, lease.expire - now)
            print(string.format("%-16s\t%-14s\t%d mins",
                    mac:sub(1,12) .. "...",
                    lease.ip,
                    math.floor(remaining / 60)))
            count = count + 1
        end
    end

    if count == 0 then
        print("No active leases")
    end
    print("============================\n")
end

local function showStats()
    local uptime = computer.uptime() - statistics.start_time
    local hours = math.floor(uptime / 3600)
    local minutes = math.floor((uptime % 3600) / 60)
    local seconds = math.floor(uptime % 60)

    print("\n===== Statistics (" .. mode .. " mode) =====")
    print("Network Card : " .. modem.address)
    print("Uptime       : " .. string.format("%02d:%02d:%02d", hours, minutes, seconds))
    print("DHCP Requests: " .. statistics.dhcp_requests)
    print("DHCP Offers  : " .. statistics.dhcp_offers)
    print("DHCP ACKs    : " .. statistics.dhcp_acks)
    print("ARP Requests : " .. statistics.arp_requests)
    print("ARP Replies  : " .. statistics.arp_replies)
    print("ICMP Requests: " .. statistics.icmp_requests)
    print("ICMP Replies : " .. statistics.icmp_replies)
    print("============================\n")
end

local function handleCommand(input)
    local parts = {}
    for word in input:gmatch("%S+") do
        table.insert(parts, word:lower())
    end

    if #parts == 0 then return true end

    local cmd = parts[1]

    if cmd == "help" then
        showHelp()

    elseif cmd == "ipconfig" then
        if #parts == 1 then
            showIPConfig()
        elseif mode == "client" then
            if parts[2] == "/release" then
                print("Releasing DHCP lease...")
                networkConfig.ip = nil
                networkConfig.lease_time = 0
                saveConfig()
                print("IP address released")
            elseif parts[2] == "/renew" then
                if networkConfig.dhcp_enabled then
                    print("Renewing DHCP lease...")
                    performDHCP()
                else
                    print("DHCP is disabled")
                end
            elseif parts[2] == "/set" and parts[3] then
                networkConfig.ip = parts[3]
                networkConfig.dhcp_enabled = false
                networkConfig.lease_time = 0
                arpTable[networkConfig.ip] = {
                    mac = modem.address,
                    timestamp = computer.uptime()
                }
                saveConfig()
                print("Static IP set to: " .. parts[3])
            end
        else
            print("ipconfig subcommands only available in client mode")
        end

    elseif cmd == "ping" then
        if parts[2] then
            local count = tonumber(parts[3]) or 4
            ping(parts[2], count)
        else
            print("Usage: ping <target> [count]")
        end

    elseif cmd == "arp" then
        if #parts == 1 then
            showArpTable()
        elseif parts[2] then
            local mac = arpRequest(parts[2])
            if mac then
                print(parts[2] .. " is at " .. mac)
            else
                print("Could not resolve " .. parts[2])
            end
        end

    elseif cmd == "dhcp" and mode == "client" then
        if parts[2] == "enable" then
            networkConfig.dhcp_enabled = true
            saveConfig()
            print("DHCP enabled")
        elseif parts[2] == "disable" then
            networkConfig.dhcp_enabled = false
            saveConfig()
            print("DHCP disabled")
        else
            print("Usage: dhcp <enable|disable>")
        end

    elseif cmd == "list" then
        showLeases()

    elseif cmd == "kick" and mode == "server" then
        if parts[2] then
            local targetIP = parts[2]
            for mac, lease in pairs(dhcpLeases) do
                if lease.ip == targetIP then
                    dhcpLeases[mac] = nil
                    arpTable[targetIP] = nil
                    print("Kicked client with IP: " .. targetIP)
                    return true
                end
            end
            print("Client with IP " .. targetIP .. " not found")
        else
            print("Usage: kick <ip>")
        end

    elseif cmd == "stats" then
        showStats()

    elseif cmd == "save" then
        if saveConfig() then
            print("Configuration saved")
        else
            print("Failed to save configuration")
        end

    elseif cmd == "reload" then
        if loadConfig() then
            print("Configuration reloaded")
        else
            print("Failed to load configuration")
        end

    elseif cmd == "exit" then
        return false

    else
        print("Unknown command: " .. cmd .. " (type 'help' for available commands)")
    end

    return true
end

-- ===============================
-- 主程序
-- ===============================

local function main(args)
    -- 解析命令行参数
    if not args[1] or (args[1] ~= "-server" and args[1] ~= "-client") then
        print("Usage: dhcp -server   or   dhcp -client")
        return
    end

    mode = args[1]:sub(2) -- 去掉"-"前缀

    -- 初始化网卡
    modem = selectNetworkCard()
    if not modem then
        print("Failed to initialize network card")
        return
    end

    -- 打开所需端口
    if mode == "server" then
        modem.open(DHCP_SERVER_PORT)
    else
        modem.open(DHCP_CLIENT_PORT)
    end
    modem.open(ARP_PORT)
    modem.open(ICMP_PORT)

    -- 加载配置
    loadConfig()

    -- 初始化模式特定的配置
    statistics.start_time = computer.uptime()

    if mode == "server" then
        dhcpLeases["server"] = {ip = serverConfig.server_ip, expire = math.huge}
        arpTable[serverConfig.server_ip] = {mac = modem.address, timestamp = computer.uptime()}
        print("DHCP Server Started")
        print("Server IP: " .. serverConfig.server_ip)
    else
        -- 添加自己的IP到ARP表(如果有的话)
        if networkConfig.ip then
            arpTable[networkConfig.ip] = {mac = modem.address, timestamp = computer.uptime()}
        end
        print("DHCP Client Started")
    end

    print("Network card: " .. modem.address)
    print("Type 'help' for available commands\n")

    -- 启动后台网络监听线程
    table.insert(threads, thread.create(networkListener))

    -- 如果是客户端模式，启动DHCP检查线程
    if mode == "client" then
        table.insert(threads, thread.create(dhcpChecker))

        -- 如果没有有效IP且启用了DHCP，立即尝试获取
        if networkConfig.dhcp_enabled and not hasValidIP() then
            performDHCP()
        end
    end

    -- CLI主循环
    while running do
        io.write(mode .. "> ")
        local input = io.read()

        if not input then break end

        if not handleCommand(input) then
            break
        end
    end

    -- 清理
    running = false
    for _, t in ipairs(threads) do
        if t then t:kill() end
    end
    saveConfig()
    modem.close(DHCP_SERVER_PORT)
    modem.close(DHCP_CLIENT_PORT)
    modem.close(ARP_PORT)
    modem.close(ICMP_PORT)
    print("DHCP " .. mode .. " stopped")
end

-- 获取命令行参数并启动
local args = {...}
main(args)
